#!/usr/bin/env python
import urllib3

from fire import Fire

import os
import io
import time
import json
from pathlib import Path
from alibabacloud_imageenhan20190930.client import Client
from alibabacloud_imageenhan20190930.models import (
    GenerateSuperResolutionImageAdvanceRequest,
    GetAsyncJobResultRequest,
)
from alibabacloud_tea_openapi.models import Config
from alibabacloud_tea_util.models import RuntimeOptions

# 禁用SSL证书验证警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
# 创建不验证SSL的HTTP管理器
http = urllib3.PoolManager(cert_reqs="CERT_NONE", assert_hostname=False)


@Fire
def main(scale):
    config = Config(
        access_key_id=os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID"),
        access_key_secret=os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET"),
        endpoint="imageenhan.cn-shanghai.aliyuncs.com",
        region_id="cn-shanghai",
    )

    input_dir = Path("out")
    output_dir = Path("scale")
    output_dir.mkdir(exist_ok=True)

    image_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"}
    image_files = [
        f for f in input_dir.iterdir() if f.suffix.lower() in image_extensions
    ]

    if not image_files:
        print("未在out目录中找到图片文件")
        return

    print(f"找到 {len(image_files)} 个图片文件")

    client = Client(config)

    for img_file in sorted(image_files):
        output_path = output_dir / img_file.name
        if output_path.exists():
            print(f"文件 {img_file} 已处理过，跳过")
            continue
        # 读取本地图片文件
        with open(img_file, "rb") as f:
            img_data = f.read()

        # 创建超分辨率请求
        img = io.BytesIO(img_data)
        generate_super_resolution_image_request = (
            GenerateSuperResolutionImageAdvanceRequest(
                scale=scale,
                output_format="jpg",
                output_quality=100,
                image_url_object=img,
            )
        )

        runtime = RuntimeOptions()

        while True:
            try:
                response = client.generate_super_resolution_image_advance(
                    generate_super_resolution_image_request, runtime
                )
                job_id = response.body.request_id
                print(f"任务提交成功，任务ID: {job_id}")
                break
            except Exception as e:
                print(f"任务提交失败：{e}")

        # 等待任务完成（轮询查询结果）
        max_wait_time = 300  # 最大等待时间5分钟
        wait_time = 0
        interval = 5  # 每5秒查询一次

        while wait_time < max_wait_time:
            time.sleep(interval)
            wait_time += interval

            try:
                result_response = client.get_async_job_result(
                    GetAsyncJobResultRequest(job_id=job_id)
                )
            except Exception as e:
                print(f"查询结果失败：{e}")
                continue

            response_map = result_response.body.to_map()
            data = response_map.get("Data") or response_map.get("data")

            status = None
            if isinstance(data, dict):
                status = data.get("Status") or data.get("status")

            if status == "SUCCEEDED" or status == "PROCESS_SUCCESS":
                if not (data and isinstance(data, dict)):
                    print("任务成功但响应数据格式不正确")
                    break

                result_str = data.get("Result") or data.get("result")
                if not (result_str and isinstance(result_str, str)):
                    print("任务成功但无法获取结果字符串")
                    break

                result_json = json.loads(result_str)
                image_url = result_json.get("resultUrl") or result_json.get("ResultUrl")
                if not image_url:
                    print("任务成功但无法获取图片URL")
                    break

                response = http.request("GET", image_url)
                image_bytes = response.data

                # 构建输出路径，保持相同的文件名

                # 保存处理后的图片
                with open(output_path, "wb") as output_file:
                    output_file.write(image_bytes)

                print(f"已保存处理后的图片到 {output_path}")
                break
            elif status == "FAILED":
                print(f"任务失败: {response_map.get('Message', 'Unknown error')}")
                break
            elif status == "PROCESSING":
                print(f"任务处理中，继续等待... ({wait_time}s/{max_wait_time}s)")
                continue
            else:
                print(f"未知状态: {status}")
                break
        else:
            print(f"任务等待超时 ({max_wait_time}秒)")

    print("所有图片处理完成")
